import pandas as pd

def merge_device_csvs(cpu_file, dpu_file, output_file, convert_dpu_to_cpu=False):
    """
    合并包含不同device_type的CSV文件
    
    参数:
        cpu_file (str): 包含device_type=cpu的CSV文件路径
        dpu_file (str): 包含device_type=dpu的CSV文件路径
        output_file (str): 合并后的输出文件路径
        convert_dpu_to_cpu (bool): 是否将dpu转换为cpu (默认为False)
    """
    try:
        # 读取两个CSV文件
        df_cpu = pd.read_csv(cpu_file)
        df_dpu = pd.read_csv(dpu_file)
        
        # 检查列是否一致
        if list(df_cpu.columns) != list(df_dpu.columns):
            print("错误：两个CSV文件的列不匹配")
            return
        
        # 可选：将dpu转换为cpu
        if convert_dpu_to_cpu:
            df_dpu['device_type'] = 'cpu'
        
        # 合并数据
        merged_df = pd.concat([df_cpu, df_dpu], ignore_index=True)
        
        # 保存结果
        merged_df.to_csv(output_file, index=False)
        print(f"成功合并文件，结果保存到: {output_file}")
        print(f"总记录数: {len(merged_df)} (CPU: {len(df_cpu)}, DPU: {len(df_dpu)})")
        
    except Exception as e:
        print(f"处理过程中发生错误: {str(e)}")

if __name__ == "__main__":
    import argparse
    
    # 设置命令行参数
    parser = argparse.ArgumentParser(description='合并包含不同device_type的CSV文件')
    parser.add_argument('cpu_file', help='包含device_type=cpu的CSV文件路径')
    parser.add_argument('dpu_file', help='包含device_type=dpu的CSV文件路径')
    parser.add_argument('output_file', help='合并后的输出文件路径')
    parser.add_argument('--convert', action='store_true', 
                       help='是否将dpu转换为cpu', default=False)
    
    args = parser.parse_args()
    
    # 执行合并
    merge_device_csvs(args.cpu_file, args.dpu_file, args.output_file, args.convert)